from re import T
import torch
import torch.nn as nn
import torch.nn.functional as F

from models.osmnet import L2NormDense
from models.custom import DeepResDDFN6b, DeepRRDDFN6

class FPN_top2bottom(nn.Module):

    def __init__(self, channels=[32, 64, 128]):
        super().__init__()
        self.channels = channels

        ch = max(channels)
        ch_sum = 128
        self.fusion = []
        for i in range(1, -1, -1):
            ch_sum += self.channels[i]
            self.add_module(
                f"fusion{i}", 
                nn.Conv2d(ch_sum, ch, kernel_size=1, bias=False)
            )
            self.fusion.append(getattr(self, f"fusion{i}"))
            ch_sum = ch

    def forward(self, x):
        
        ft = []
        net = iter(self.fusion)
        for f_map in reversed(x):

            ft.append(f_map)
            if len(ft) == 2:
                ft_cat = torch.cat(ft, dim=1)
                cur_fusion = next(net)
                out = cur_fusion(ft_cat)
                ft = [out]

        return ft[0]


class FPNtbRRDDFN6(DeepRRDDFN6):

    def __init__(self, out_ch=9):
        super().__init__(out_ch, ret_ip=True)
        self.fpn = FPN_top2bottom()

    def forward(self, x):
        fts = super().forward(x)
        ms_fts = self.fpn(fts)
        desc = self.reduce(ms_fts)

        norm_desc = L2NormDense()(desc)

        return desc, norm_desc


class FPNtbResDDFN6b(DeepResDDFN6b):

    def __init__(self, out_ch=9):
        super().__init__(out_ch, ret_ip=True)
        self.fpn = FPN_top2bottom()

    def forward(self, x):
        fts = super().forward(x)
        ms_fts = self.fpn(fts)
        desc = self.reduce(ms_fts)

        norm_desc = L2NormDense()(desc)

        return desc, norm_desc


if __name__ == "__main__1":

    net = FPN_top2bottom(channels=[32, 64, 128]).cuda()

    ft = [
        torch.zeros((2, 32, 100, 100)).cuda(), 
        torch.zeros((2, 64, 100, 100)).cuda(), 
        torch.zeros((2, 128, 100, 100)).cuda()
    ]

    out = net(ft)
    print(out.shape)


if __name__ == "__main__":

    net = FPNtbResDDFN6b().cuda()

    ft = torch.zeros((2, 1, 100, 100)).cuda()

    out = net(ft)
    print(out.shape)


